#import math
import matplotlib
matplotlib.rcParams['units'] = True
import matplotlib.units as units
import matplotlib.ticker as ticker
#from matplotlib.axes import Axes
from matplotlib.cbook import iterable, is_string_like


from Physics.Units.PhysicalQuantities import PhysicalUnit as PU 
from Physics.Units.PhysicalQuantities import PhysicalQuantity as PQ
from Physics.Units.PhysicalQuantities import _unit_table,  isPhysicalUnit, isPhysicalQuantity

from numpy import isreal, isscalar 
""""
# SI unit definitions

_base_names = ['m', 'kg', 's', 'A', 'K', 'mol', 'cd', 'rad', 'sr']
"""


def addUnitsToX(axs,xunit):
    axs.xaxis.set_units(xunit)
    for line in axs.lines:
        xdata = line.get_xdata()
        if not isPhysicalQuantity(xdata[0]):
            line.set_xdata(xdata*PQ(1,xunit))
            
def addUnitsToY(axs,yunit):
    axs.yaxis.set_units(yunit)
    for line in axs.lines:
        ydata = line.get_ydata()
        if not isPhysicalQuantity(ydata[0]):
            line.set_ydata(ydata*PQ(1,yunit))
            
def addUnitsToXY(axs,xunit,yunit):
    addUnitsToX(axs,xunit)
    addUnitsToY(axs,yunit)
    
def addUnitsToLine(line,xunit,yunit):
    xdata = line.get_xdata()
    ydata = line.get_ydata()
    line.set_xdata(xdata*PQ(1,xunit))
    line.set_ydata(ydata*PQ(1,yunit))
    


    

class BasicUnitConverter(units.ConversionInterface):

    def axisinfo(unit,axis):
        #print 'return AxisInfo instance for x and unit :', unit

        if unit == _unit_table['rad']:
            return units.AxisInfo(
              majloc=ticker.MultipleLocator(base=nx.pi/2),
              majfmt=ticker.FuncFormatter(rad_fn),
              label=unit.name()
                )
        elif unit == _unit_table['deg']:
            return units.AxisInfo(
              majloc=ticker.AutoLocator(),
              majfmt=ticker.FormatStrFormatter(r'$%i^\circ$'),
              label=unit.name()
                )
        elif unit is not None:
            if isPhysicalUnit(unit):
                return units.AxisInfo(label=unit.name())
            elif isPhysicalQuantity(unit):
                return units.AxisInfo(label=unit.unit.name())
        return None
    axisinfo = staticmethod(axisinfo)
    
    
    def convert(self,val, unit,axis):
        #print 'converting...'
        if iterable(val):
                if is_string_like(val): #strings will lead to infinite loop!!
                    return val
                return [ self.convert(o,unit,axis) for o in val]
        
        if isPhysicalQuantity(val):
            return val.inUnitsOf(unit).getValue()
        
        #if isreal(val) and (unit is not None):# for doing scaling of axis
        #    print 'scaling'
        #    return val/unit
        else:
            return val
        
        #print 'convert '
        #if iterable(val):
        #print 'found iterable...' 
        #    if not isPhysicalQuantity(val[0]):
        #        return val
         #   else:
         #       newval = []
         #       for thisval in val:
         #           #print thisval
         #           newval = newval + [thisval.inUnitsOf(unit).getValue(),]
         #       return newval
        #else:
         #   if isPhysicalQuantity(val):
                #print 'found PQ...'
          #      return val.inUnitsOf(unit).getValue()
           # else:
                #print 'found non PQ...'
    #            return val
    convert = classmethod(convert)
    
    def toPQ(self,val,unit):
    
        if iterable(val):
            if is_string_like(val):
                return val
            return [self.toPQ(o,unit) for o in val]
        
        if isreal(val):
            return PQ(val,unit)
        else:
            return val
    toPQ= classmethod(toPQ)
    
    def add_xunit(axes,unit):
        axes.xaxis.set_units(unit)
        for line in axes.lines:
            xdata = line.get_xdata()
            if not isPhysicalQuantity(xdata[0]):
                line.set_xdata(xdata*PQ(1,unit))
    add_xunit = staticmethod(add_xunit)
 
        
    def add_yunit(axes,unit):
        axes.yaxis.set_units(unit)
        for line in axes.lines:
            ydata = line.get_ydata()
            if not isPhysicalQuantity(ydata[0]):
                line.set_ydata(ydata*PQ(1,unit))
    add_yunit = staticmethod(add_yunit)
        
    @staticmethod
    def default_units(x,axis):
        'return the default unit for x or None'
        if iterable(x):
            #print x
            return BasicUnitConverter.default_units(x[0], axis) #[BasicUnitConverter.default_units(val,axis) for val in x]
        
        return x.unit
#    default_units = staticmethod(default_units)

# radians formatting
def rad_fn(x,pos=None):
  n = int((x / nx.pi) * 2.0 + 0.25)
  if n == 0:
    return '0'
  elif n == 1:
    return r'$\pi/2$'
  elif n == 2:
    return r'$\pi$'
  elif n % 2 == 0:
    return r'$%s\pi$' % (n/2,)
  else:
    return r'$%s\pi/2$' % (n,)

basicConverter = BasicUnitConverter()
units.registry[PU] = basicConverter
units.registry[PQ] = basicConverter